from __future__ import division
#import ImageMetrics as IM
#import Metric
import pylab
import Image
import numpy as np

import JT
import JT.Data2 as Data
#import Calibration

class Test():
    def __init__(self):    
        """
        This function will plot a metric over the given range, for one element of 
        the pose.
        Inputs:
            eval_range: an iterable range of positions at which to calculate the metric
            eval_param: the element of the pose that will be modified over the range
        output:
            val: array of values returned by the metric at each position in eval_range
        """
#===============================================================================
#    Settings
#===============================================================================   
        # ****Make sure angle and translation ranges are the same length
        # Iterate over a range of translations from -30 to 30 mm
        numel = 10
        pos = 5
        ang = np.pi/180*pos
        
        
        self.trans_range = np.arange(-pos,pos+pos/numel,pos/numel)
        #self.trans_range = np.zeros(numel)
        # Iterate over a range of angles from -pi/4 (45deg) to pi/4 in 2.864 deg increments
        self.angle_range = np.arange(-ang,ang+ang/numel,ang/numel)
        #self.angle_range = np.zeros(numel)
        # Define which data to use for registration:
        self.Data = Data.DataStructure()
        
        self.reg=JT.MultiIterationRegistration() 
        self.reg.drr = JT.DRR()  
        volume = JT.Volume2()
        self.reg.drr.SetVolume(volume) 
        
        

    
#===============================================================================
#    Start Metric evaluation loop
#==============================================================================
    def run_test(self,xml_file=None,method='bpc',tag='',
                 frame = 1, volume = 1, save_fig=True):
        
            
#===============================================================================
#    Setup DRR
#===============================================================================      
        self.Data.load(xml_file)
        
        # Pull the Volume and FixedImage elements out of the data structure
        Volume = self.Data.Tree.Frame[frame].Volume[volume]
        FixedImage = self.Data.Tree.Frame[frame].FixedImage[0]       
        # Load all of the data elements from the structure into the reg class
        self.reg.LoadCalFile(self.Data.get_path(FixedImage.Cal))
        self.reg.drr.xrayCamera.DoCalibration()
        #self.reg.LoadStaFile(data.get_path(Volume.STA))
        self.reg.SetVolumeFileName(self.Data.get_path(Volume))
        self.reg.SetFixedImageFileName(self.Data.get_path(FixedImage))
        self.reg.SetRegionOfInterest(self.Data.get_roi(FixedImage.path,Volume.path))
        #self.reg.SetCenter(self.Data.get_center(FixedImage.path,Volume.path))
        
        starting_pose = self.Data.get_pose(FixedImage.path,Volume.path)
        self.reg.SetImageMetric(method)
        
        (threshold, level) = self.Data.get_global_display_settings()
        self.reg.drr.volume.SetOpacityLevel(level)
        self.reg.drr.volume.SetOpacityThreshold(threshold)
        #self.reg.drr.renWin.SetOffScreenRendering(1)
        self.reg.StartNewLog() 
        self.reg.settings['plot']=False
        self.reg.settings['debug']=False
        self.reg.drr.OffScreenRenderingOn()
        
        # Set the origin to be the center of the data
        self.reg.drr.volume.SetOrientation(0,0,0)
        center = np.asarray(self.reg.drr.volume.GetCenter()) -\
                 np.asarray(self.reg.drr.volume.GetPosition())
        self.reg.drr.volume.SetOrigin(center)
        
        # Offset the poses by the optional starting pose
        all_poses = np.array([self.trans_range + starting_pose[0],
                              self.trans_range + starting_pose[1],
                              self.trans_range + starting_pose[2],
                              self.angle_range + starting_pose[3],
                              self.angle_range + starting_pose[4],
                              self.angle_range + starting_pose[5]]).transpose()
        all_poses_angles = all_poses.copy()
        all_poses_angles[:,3:6] = all_poses_angles[:,3:6]*180/np.pi
        starting_pose_degrees = starting_pose.copy()
        starting_pose_degrees[3:6] = starting_pose[3:6]*180/np.pi
        best_pose = np.zeros(6)
        
        num_evals = len(all_poses)
        val = np.zeros((num_evals,6))
        starting_val = 0
        pylab.ion()
        fig = pylab.figure()
        fig.canvas.setFixedSize(800,900)
        fig.set_size_inches(8,10)
        fig.subplots_adjust(wspace=0.35,hspace=0.4)
        pylab.suptitle(tag,fontsize=14,fontweight='bold')
        ax = [range(6)]
        
        for dim in range(6):
            min_val = 1e300
            # loop through each dimension of the pose (Tx, Ty, Tz, Rx, Ry, Rz)
            # only want to vary one parameter at a time, all others are set to 0 using a mask
            pose_mask = np.ones((num_evals,6))*starting_pose
            pose_mask[:,dim]=all_poses[:,dim]
            poses = pose_mask
            starting_val = self.reg.GetValue(starting_pose)
            for i,pose in enumerate( poses ):
                val[i,dim] = self.reg.GetValue(pose)
                if val[i,dim] < min_val:
                    best_pose[dim] = pose[dim]
                    min_val = val[i,dim]
                print "    Metric: %s, Iteration: %i, Metric Value: %f" % \
                      (self.reg.metric_settings['method'],i,val[i,dim])
 
            #plot results
            # Create a subplot for each dimension
            if dim<3:
                ax = fig.add_subplot(3,2,dim*2+1)
            else:
                ax = fig.add_subplot(3,2,(dim-3)*2+2)
            ax.plot(all_poses_angles[:,dim],val[:,dim])
            ax.plot([starting_pose_degrees[dim]],[starting_val],color='#FF4500',marker='o')
            ax.set_title(self.reg.metric_settings['method']+' Dim:'+str(dim))
            ax.set_xlabel('Pose')
            ax.set_ylabel('Metric Value')
            ax.set_xlim(all_poses_angles[0,dim],all_poses_angles[-1,dim])
            pylab.draw()
            
        
        resultsFolder = 'C:/Users/bryan/bryan-code/trunk/Results/'
        save_name = resultsFolder+method.upper()+'_'+tag+'_'+'_metric_function.png'
        pylab.draw()
        best_pose[3:6] = best_pose[3:6]*180/np.pi
        print "It looks like %s might be a good pose to try." % str(best_pose)
        print best_pose[0:3], best_pose[3:6]*np.pi/180
        #pylab.show()
        if save_fig:
            pylab.savefig(save_name)
        

if __name__ == "__main__":
    #method_list = ['nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc', 'gd', 'gc']
    method_list = ['bpc']
    file = r'C:\Users\bryan\bryan-code\trunk\Tests\test_relpath.xml'
    frame = [0,1]
    proj = [0,1]
    body = [[0,1,2],[0,1,2,3,4]]
    #body=[[1],[1]]
    if 1:
        f=1
        v=1
        
        #pose=1
        tag = 'RealData-Frame-%s Volume-%s' % (str(f),str(v))
        method = 'bpc'
        test = Test()
        test.run_test(xml_file=file,method=method,tag=tag,frame=f,volume=v,save_fig=True)
    else:
        for v in frame:
            for p in proj:
                for b in body[v]:
                    # to label the resulting image that will be saved.
                    tag = 'frame-%s Proj-%s Body-%s' % (str(v),str(p),str(b))
                    test = Test()
                    for method in method_list:
                        test.run_test(method,tag,frame=v,proj=p,body=b,save_fig=True)
                    test = None


        
        